from __future__ import annotations  # 延迟评估注解(类型声明)

import math
from typing import List, Union

import matplotlib.pyplot as plt
import numpy as np


class Node:
    def __init__(self, grid_pos: tuple[int, int], g: float, h: float, parent: Node = None):
        self.grid_pos = tuple(grid_pos)
        self.g = g
        self.h = h
        self.parent = parent
        self.f = self.g + self.h


class JPS:
    def __init__(self, width: int, height: int, map_data: Union[List[List[int]], np.ndarray] = None):
        self.start_grid_pos = (0, 0)
        self.goal_grid_pos = (0, 0)
        self.width = width
        self.height = height
        self.map_data = map_data

        self.open = {}  # pos:node
        self.close = {}

        self.motion_directions = [[1, 0], [0, 1], [0, -1], [-1, 0], [1, 1], [1, -1], [-1, 1], [-1, -1]]

    def run(self, start_grid_pos: tuple[int, int], goal_grid_pos: tuple[int, int]):
        self.start_grid_pos = start_grid_pos
        self.goal_grid_pos = tuple(goal_grid_pos)
        start_node = Node(start_grid_pos, 0, self.getH(self.start_grid_pos, self.goal_grid_pos))
        self.open[start_node.grid_pos] = start_node

        while True:

            # 寻路失败，结束循环
            if not self.open:
                print("未找到路径")
                break

            current_node = min(self.open.values(), key=lambda x: x.f)  # f值最小的节点

            # 找到路径， 返回结果
            if current_node.grid_pos == self.goal_grid_pos:
                path = self.findPath(current_node)
                new_path = self.simplify_path(path)  # 简化路径
                print("找到路径！")
                return new_path

            # 扩展节点
            self.extendNode(current_node)

            # 更新节点
            self.close[current_node.grid_pos] = current_node
            del self.open[current_node.grid_pos]

    def extendNode(self, current_node: Node):
        """
        根据当前节点，扩展节点（只有跳点才可以扩展）
        input
        ----------
        current_node: 当前节点对象
        """
        neighbours = self.getPruneNeighbours(current_node)
        for n in neighbours:
            jp = self.getJumpNode(n, current_node.grid_pos)  # 跳点
            if jp:
                if jp in self.close:
                    continue

                new_node = Node(jp, current_node.g + self.getG(jp, current_node.grid_pos),
                                self.getH(jp, self.goal_grid_pos), current_node)
                if jp in self.open:
                    if new_node.f < self.open[jp].f:
                        self.open[jp].parent = current_node
                        self.open[jp].f = new_node.f
                else:
                    self.open[jp] = new_node

    def getJumpNode(self, now: tuple[int, int], pre: tuple[int, int]) -> Union[tuple[int, int], None]:
        """
        计算跳点
        input
        ----------
        now: 当前节点坐标
        pre: 上一节点坐标
        output
        ----------
        若有跳点，返回跳点坐标；若无跳点，则返回None
        """
        x_direction = int((now[0] - pre[0]) / abs(now[0] - pre[0])) if now[0] - pre[0] != 0 else 0
        y_direction = int((now[1] - pre[1]) / abs(now[1] - pre[1])) if now[1] - pre[1] != 0 else 0

        if now == self.goal_grid_pos:  # 如果当前节点是终点，则为跳点（条件1）
            return now

        if self.hasForceNeighbours(now, pre):  # 如果当前节点包含强迫邻居，则为跳点（条件2）
            return now

        if abs(x_direction) + abs(y_direction) == 2:  # 若为斜线移动，则判断水平和垂直方向是否有满足上述条件1和2的点，若有，则为跳点（条件3）

            if (self.getJumpNode((now[0] + x_direction, now[1]), now) or
                    self.getJumpNode((now[0], now[1] + y_direction), now)):
                return now

        if self.isPass(now[0] + x_direction, now[1] + y_direction):  # 若当前节点未找到跳点，朝当前方向前进一步，继续寻找跳点，直至不可达
            jp = self.getJumpNode((now[0] + x_direction, now[1] + y_direction), now)
            if jp:
                return jp

        return None

    def hasForceNeighbours(self, now: tuple[int, int], pre: tuple[int, int]) -> bool:
        """
        根据当前节点坐标和上一节点坐标判断当前节点是否拥有强迫邻居
        input
        ----------
        now: 当前坐标
        pre: 上一坐标
        output
        ----------
        若拥有强迫邻居，则返回True；否则返回False
        """
        x_direction = now[0] - pre[0]
        y_direction = now[1] - pre[1]

        # 若为直线移动
        if abs(x_direction) + abs(y_direction) == 1:

            if abs(x_direction) == 1:  # 水平移动
                if (self.isPass(now[0] + x_direction, now[1] + 1) and not self.isPass(now[0], now[1] + 1)) or \
                        (self.isPass(now[0] + x_direction, now[1] - 1) and not self.isPass(now[0], now[1] - 1)):
                    return True
                else:
                    return False

            elif abs(y_direction) == 1:  # 垂直移动

                if (self.isPass(now[0] + 1, now[1] + y_direction) and not self.isPass(now[0] + 1, now[1])) or \
                        (self.isPass(now[0] - 1, now[1] + y_direction) and not self.isPass(now[0] - 1, now[1])):
                    return True
                else:
                    return False

            else:
                raise Exception("错误，直线移动中只能水平或垂直移动！")

        # 若为斜线移动
        elif abs(x_direction) + abs(y_direction) == 2:
            if (self.isPass(now[0] + x_direction, now[1] - y_direction) and not self.isPass(now[0],
                                                                                            now[1] - y_direction)) or \
                    (self.isPass(now[0] - x_direction, now[1] + y_direction) and not self.isPass(now[0] - x_direction,
                                                                                                 now[1])):
                return True
            else:
                return False

        else:
            raise Exception("错误，只能直线移动或斜线移动！")

    def getH(self, current: tuple[int, int], goal: tuple[int, int], func: str = "Euclidean") -> float:
        """
        根据当前坐标和终点坐标计算启发值，包含：
            曼哈顿距离（Manhattan Distance）
            欧几里得距离（Euclidean Distance）  默认
            切比雪夫距离（Chebyshev Distance）
        input
        ----------
        current: 当前节点坐标
        goal: 目标节点坐标
        func: 启发函数，默认为欧几里得距离
        """
        current_x = current[0]
        current_y = current[1]
        goal_x = goal[0]
        goal_y = goal[1]

        if func == "Manhattan":  # 适用于格点地图上，移动限制为上下左右四个方向的情况
            h = abs(current_x - goal_x) + abs(current_y - goal_y)

        elif func == "Euclidean":  # 适用于可以自由移动至任何方向的情况，例如在平面上自由移动
            h = math.hypot(current_x - goal_x, current_y - goal_y)

        elif func == "Chebyshev":  # 适用于八方向移动的格点地图（上下左右及对角线）
            h = max(abs(current_x - goal_x), abs(current_y - goal_y))

        else:
            raise Exception("错误，不支持该启发函数。目前支持：Manhattan、Euclidean（默认）、Chebyshev。")

        return h

    def getG(self, pos1: tuple[int, int], pos2: tuple[int, int]) -> float:
        """
        根据两坐标计算代价值（直走或斜走，直接计算欧几里得距离就可）
        input
        ----------
        pos1: 坐标1
        pos2: 坐标2
        output
        ----------
        返回代价值
        """
        return math.hypot(pos1[0] - pos2[0], pos1[1] - pos2[1])

    def getPruneNeighbours(self, current_node: Node) -> list[tuple[int, int]]:
        """
        获得裁剪之后的邻居节点坐标列表（包含自然邻居和强迫邻居，其相关概念参考：https://fallingxun.github.io/post/algorithm/algorithm_jps/）
        input
        ----------
        current_node: 当前节点对象
        output
        ----------
        prune_neighbours: 裁剪邻居坐标列表
        """
        prune_neighbours = []

        if current_node.parent:
            motion_x = int((current_node.grid_pos[0] - current_node.parent.grid_pos[0]) / abs(
                current_node.grid_pos[0] - current_node.parent.grid_pos[0])) if current_node.grid_pos[0] - \
                                                                                current_node.parent.grid_pos[
                                                                                    0] != 0 else 0  # 方向不分大小，所以要除以长度
            motion_y = int((current_node.grid_pos[1] - current_node.parent.grid_pos[1]) / abs(
                current_node.grid_pos[1] - current_node.parent.grid_pos[1])) if current_node.grid_pos[1] - \
                                                                                current_node.parent.grid_pos[
                                                                                    1] != 0 else 0  # 方向不分大小，所以要除以长度

            if abs(motion_x) + abs(motion_y) == 1:  # 直线

                # 自然邻居
                if self.isPass(current_node.grid_pos[0] + motion_x, current_node.grid_pos[1] + motion_y):
                    prune_neighbours.append((current_node.grid_pos[0] + motion_x, current_node.grid_pos[1] + motion_y))

                # 强迫邻居
                if abs(motion_x) == 0:  # 垂直走
                    if not self.isPass(current_node.grid_pos[0] + 1, current_node.grid_pos[1]) and self.isPass(
                            current_node.grid_pos[0] + 1, current_node.grid_pos[1] + motion_y):
                        prune_neighbours.append((current_node.grid_pos[0] + 1, current_node.grid_pos[1] + motion_y))
                    if not self.isPass(current_node.grid_pos[0] - 1, current_node.grid_pos[1]) and self.isPass(
                            current_node.grid_pos[0] - 1, current_node.grid_pos[1] + motion_y):
                        prune_neighbours.append((current_node.grid_pos[0] - 1, current_node.grid_pos[1] + motion_y))
                else:  # 水平走
                    if not self.isPass(current_node.grid_pos[0], current_node.grid_pos[1] + 1) and self.isPass(
                            current_node.grid_pos[0] + motion_x, current_node.grid_pos[1] + 1):
                        prune_neighbours.append((current_node.grid_pos[0] + motion_x, current_node.grid_pos[1] + 1))
                    if not self.isPass(current_node.grid_pos[0], current_node.grid_pos[1] - 1) and self.isPass(
                            current_node.grid_pos[0] + motion_x, current_node.grid_pos[1] - 1):
                        prune_neighbours.append((current_node.grid_pos[0] + motion_x, current_node.grid_pos[1] - 1))

            elif abs(motion_x) + abs(motion_y) == 2:  # 对角线

                # 自然邻居
                if self.isPass(current_node.grid_pos[0] + motion_x, current_node.grid_pos[1] + motion_y):
                    prune_neighbours.append((current_node.grid_pos[0] + motion_x, current_node.grid_pos[1] + motion_y))
                if self.isPass(current_node.grid_pos[0] + motion_x, current_node.grid_pos[1]):
                    prune_neighbours.append((current_node.grid_pos[0] + motion_x, current_node.grid_pos[1]))
                if self.isPass(current_node.grid_pos[0], current_node.grid_pos[1] + motion_y):
                    prune_neighbours.append((current_node.grid_pos[0], current_node.grid_pos[1] + motion_y))

                # 强迫邻居
                if not self.isPass(current_node.grid_pos[0] - motion_x, current_node.grid_pos[1]) and self.isPass(
                        current_node.grid_pos[0] - motion_x, current_node.grid_pos[1] + motion_y):
                    prune_neighbours.append((current_node.grid_pos[0] - motion_x, current_node.grid_pos[1] + motion_y))
                if not self.isPass(current_node.grid_pos[0], current_node.grid_pos[1] - motion_y) and self.isPass(
                        current_node.grid_pos[0] + motion_x, current_node.grid_pos[1] - motion_y):
                    prune_neighbours.append((current_node.grid_pos[0] + motion_x, current_node.grid_pos[1] - motion_y))

            else:
                raise Exception("错误，只能对角线和直线行走！")

        else:
            for dir in self.motion_directions:
                if self.isPass(current_node.grid_pos[0] + dir[0], current_node.grid_pos[1] + dir[1]):
                    prune_neighbours.append((current_node.grid_pos[0] + dir[0], current_node.grid_pos[1] + dir[1]))

        return prune_neighbours

    def isPass(self, grid_x: int, grid_y: int) -> bool:
        """
        判断该栅格坐标是否可以通过
        input
        ----------
        grid_x: 栅格x坐标
        grid_y: 栅格y坐标
        output
        ----------
        若可以通过，为True，反之为False
        """
        if 0 <= grid_x < self.width and 0 <= grid_y < self.height:
            if self.map_data[grid_x][grid_y] != True or [grid_x, grid_y] == self.goal_grid_pos:
                return True
            else:
                return False
        else:
            return False

    def findPath(self, node: Node) -> list[list[int]]:
        """
        根据给定节点回溯到开始点，找到路径
        """
        path_x = [node.grid_pos[0]]
        path_y = [node.grid_pos[1]]
        while node.parent:
            node = node.parent
            path_x.append(node.grid_pos[0])
            path_y.append(node.grid_pos[1])

        return [path_x[::-1], path_y[::-1]]

    def path_safe(self, pt1, pt2):
        l1_2 = np.array(pt2) - np.array(pt1)
        if np.linalg.norm(l1_2) < 1e-6:  # 两点重合
            return True

        l1_2 = l1_2 / np.linalg.norm(l1_2)  # 单位向量

        pt_test = np.array(pt1)
        while True:
            if self.map_data[int(pt_test[0]), int(pt_test[1])]:
                return False
            pt_test = pt_test + l1_2  # 沿着方向前进
            if pt_test[0] > pt2[0] or pt_test[1] > pt2[1]:
                break

        return True

    def simplify_path(self, path: list[list[int]]):
        new_path = [[], []]
        for i in range(len(path[0])):
            if i <= 1:
                new_path[0].append(path[0][i])
                new_path[1].append(path[1][i])
            elif i == len(path) - 1:
                new_path[0].append(path[0][i])
                new_path[1].append(path[1][i])
            else:
                # 连线上未经过障碍物
                if self.path_safe(
                        (new_path[0][-2], new_path[1][-2]),
                        (path[0][i], path[1][i])
                ):
                    new_path[0].pop(-1)
                    new_path[1].pop(-1)
                new_path[0].append(path[0][i])
                new_path[1].append(path[1][i])
        return new_path


def path_plan(width, height, map, start, goal, grid_resolution, time_start, time_goal):
    jps = JPS(width, height, map)
    path = jps.run(start, goal)
    if path is None:
        print(f"{time_start}: {start}到{goal}的路径规划失败")
        path_out = [{
            'position': [goal[0] * grid_resolution, goal[1] * grid_resolution, 0.0],
            'arrival_time': time_goal
        }]
        safe_out = False
        return path_out, safe_out
    # 按格式输出
    path_out = []
    dist_list = []
    for i in np.arange(1, len(path[0])):
        path_out.append({
            'position': [path[0][i] * grid_resolution, path[1][i] * grid_resolution, 0.0],
            'arrival_time': None
        })
        dist_list.append(np.linalg.norm((
            (path[0][i] - path[0][i - 1]) * grid_resolution,
            (path[1][i] - path[1][i - 1]) * grid_resolution
        )))
    dist_sum = sum(dist_list)
    if dist_sum <= 1e-10:
        time_ratio = 0
    else:
        time_ratio = (time_goal - time_start) / sum(dist_list)
    for i in range(len(path_out)):
        path_out[i]['arrival_time'] = time_start + time_ratio * sum(dist_list[0:i + 1])
    safe_out = True
    print(f"{time_start}: {start}到{goal}的路径规划成功")
    return path_out, safe_out


def main():
    width = 500
    height = 700
    start = (5, 5)
    goal = (480, 690)
    map_test = np.zeros((width, height), dtype=int)
    map_test[10:200, 10:200] = 1
    for _ in range(32):
        # 随机目标位置
        tar_pos = np.random.uniform(0, 1, (1, 2)) * [width, height]  # 转换到地图坐标
        # 生成每个格子的坐标
        x, y = np.indices((width, height))  # y.shape = (h, w), x.shape = (h, w)
        # 计算到 tar_pos 的欧几里得距离
        dist = np.sqrt((x - tar_pos[0, 0]) ** 2 + (y - tar_pos[0, 1]) ** 2)
        # 距离 <= 30 的位置设为 1
        map_test[dist <= 30] = 1

    # jps = JPS(width, height, map_test)
    # path = jps.run(start, goal)
    grid_resolution = 1000
    path_res, flag = path_plan(width, height, map_test, start, goal, grid_resolution, 10.0, 100)
    print("路径规划结果：", path_res)

    fig = plt.figure()

    ax = fig.add_subplot(111)
    ax.imshow(map_test, cmap='binary', vmin=0, vmax=1)

    ax.scatter(start[1], start[0])  # 绘制图像的坐标轴与np相反
    ax.scatter(goal[1], goal[0])

    px, py = [start[0]], [start[1]]
    for i in path_res:
        px.append(i['position'][0] / grid_resolution)
        py.append(i['position'][1] / grid_resolution)
    ax.plot(py, px, '.-')
    plt.show()


if __name__ == '__main__':
    main()
